import pandas as pd
import re
import numpy as np
import folium
import patsy
import statsmodels.api as sm


def get_data():
    """
    数据获取，并获取目标数据
    :return:
    """
    df = pd.read_csv('./data/magic.csv')
    return df[['listingtype_value', 'pricelarge_value_prices', 'propertyinfo_value', 'routable_link/_text']]


def check_data(df):
    """
    查看数据概况
    :param df:
    :return:
    """
    print(df.head().T)

    print(len(df))


def prepare(df):
    """
    数据准备
    :param df:
    :return:
    """
    df = df.dropna()
    print(len(df))

    # 关键指标：listingtype_value 区分房型，并打印不同房型数量
    mu = df[df['listingtype_value'].str.contains('Apartments For Rent')]
    signal = df[df['listingtype_value'].str.contains('Apartment For')]
    print('mu: {}, signal: {}'.format(len(mu), len(signal)))

    # 目标数据：
    # 卧室数(propertyinfo_value)，浴室数(propertyinfo_value)，
    # 面积(propertyinfo_value)，地址，价格（pricelarge_value_prices）
    # 位置楼层等信息（routable_link/_text）

    # 数据质量检查
    # 检查空值
    print(df['pricelarge_value_prices'].unique())
    print(len(df['pricelarge_value_prices']==np.nan))

    print(len(df['propertyinfo_value']==np.nan))

    # 检查包含数据
    print(len(signal[~(signal['propertyinfo_value'].str.contains('Studio') | signal['propertyinfo_value'].str.contains('bd'))]))
    print(len(signal[~(signal['propertyinfo_value'].str.contains('ba'))]))

    return df


# 数据处理
def parse_info(row):
    """
    从propertyinfo_value中获取卧室、浴室和浴室面积信息
    :param row:
    :return:
    """
    row_arr = row.split('•')
    row_arr = [data.strip() for data in row_arr]
    row_dict = dict()
    for arr in row_arr:
        if 'sqft' in arr or 'bd' in arr or 'ba' in arr or 'bds' in arr:
            v, k = arr.split(' ')[: 2]
            k = 'bd' if k == 'bds' else k
            row_dict[k] = v
        elif 'Studio' in arr:
            row_dict['bd'] = 'Studio'

    return pd.Series(row_dict)


def parse_add(row):
    """
    从routable_link/_text中解析地址
    :param row:
    :return:
    """
    zip = re.search(', NY(\d+)', row)
    flr = re.search('(?:APT|#)\s+(\d+)[A-Z]+', row)
    if zip:
        zipc = zip.group(1)
    else:
        zipc = np.nan
    if flr:
        flrc = flr.group(1)
    else:
        flrc = np.nan
    return pd.Series({'Zip': zipc, 'Floor': flrc})


def pre_process(df):
    """
    数据预处理
    :param df:
    :return:
    """
    # 房屋卧室等信息
    info = df['propertyinfo_value'].apply(parse_info)
    # 位置等信息
    locat = df['routable_link/_text'].apply(parse_add)

    df_new = pd.concat([df, info, locat], axis=1)

    # 目标数据
    df_target = df_new[['pricelarge_value_prices', 'bd', 'ba', 'sqft', 'Floor', 'Zip']]
    # 租金，卧室数，浴室数，面积，楼层，邮编
    df_target.columns = ['Rent', 'Beds', 'Baths', 'Sqft', 'Floor', 'Zip']
    return df_target


def analysis(df):
    """
    数据分析
    :param df:
    :return:
    """
    print(df.describe())
    # 将 Beds 的Studio(工作室) 替换为0
    df['Beds'] = df['Beds'].apply(lambda x: 0 if x == 'Studio' else int(x))
    print(df.head())
    print(df.info())
    df['Rent'] = df['Rent'].astype(int)
    df['Baths'] = df['Baths'].astype(float)
    df['Sqft'] = df['Sqft'].str.replace(',', '')
    df['Sqft'] = df['Sqft'].astype(float)
    df['Floor'] = df['Floor'].astype(float)
    print(df.info())
    # 剔除1107层的房屋
    df = df.drop([470])
    # print(df)
    # df.to_csv('./data/test.csv')
    print(df.describe())

    # 透视表：通过邮编+卧室数来看其与价格的关系
    print(df.pivot_table('Rent', 'Beds', 'Zip', aggfunc='mean').T)
    print(df.pivot_table('Rent', 'Beds', 'Zip', aggfunc='count').T)

    # 关注工作室+1室
    # df_plot = df[df['Beds'] < 2]
    # map = folium.Map(location=[40.748817, -73.985428], zoom_start=13)
    # geo_path = './data/NY.json'
    # folium.GeoJson(geo_path, data=df_plot,
    #              columns=['Zip', 'Rent'], key_on='feature.properties.postalCode',
    #              threshold_scale=[1700, 1900, 2100, 2300, 2500, 2750],
    #              fill_color='YIOrRd', fill_opacity=0.7, line_opacity=0.2,
    #              legend_name='租金 (%)', reset=True)
    # map.create_map(path='./data/nyc.html')
    return df


def model(df):
    """
    建模
    :param df:
    :return:
    """
    # 1、建模
    func = 'Rent ~ Beds + Zip'
    y, x = patsy.dmatrices(func, df, return_type='dataframe')

    results = sm.OLS(y, x).fit()
    print(results.summary())

    # 2、预测
    predict_data = np.zeros(len(x.iloc[0]))

    # 截距
    predict_data[0] = 1
    # Zip[T.10009]
    predict_data[6] = 1
    # 卧室数
    predict_data[-1] = 2

    pre = results.predict(predict_data)
    print(pre)


def run():
    """
    主程序入口
    :return:
    """
    # 数据获取
    df = get_data()

    # 数据准备
    df_new = prepare(df)

    # 数据预处理
    df_target = pre_process(df_new)

    # 数据分析
    df_model = analysis(df_target)

    # 建模
    model(df_model)


if __name__ == '__main__':
    run()
